import numpy as np
import torch as pt
import torchvision as ptv
from torch.utils.data import DataLoader, Dataset, Subset
from torchvision.datasets import MNIST
import os
import sys
import matplotlib.pyplot as plt

np.random.seed(777)
pt.manual_seed(777)

BATCH_SIZE = 16
DATA_DIR = '../../../../../large_data/DL2/pt/cifar10'
mnist_train = ptv.datasets.CIFAR10(DATA_DIR, train=True,
                                 transform=ptv.transforms.ToTensor(),
                                 download=False)
mnist_train_transf = ptv.datasets.CIFAR10(DATA_DIR, train=True,
                                        transform=ptv.transforms.Compose([
                                            # ptv.transforms.RandomHorizontalFlip(p=1.0),
                                            # ptv.transforms.RandomVerticalFlip(p=1.0),
                                            # ptv.transforms.RandomRotation(degrees=45),
                                            ptv.transforms.RandomRotation(degrees=(45, 45)),
                                            ptv.transforms.Resize((40, 40)),
                                            ptv.transforms.RandomCrop((28, 28)),
                                            # ptv.transforms.RandomGrayscale(p=1.0),
                                            ptv.transforms.ToTensor(),
                                            ptv.transforms.Normalize((0., 0., 0.), (1., 1., 1.))
                                        ]),
                                        download=False)
dl = DataLoader(mnist_train, batch_size=BATCH_SIZE, shuffle=False)
dl_transf = DataLoader(mnist_train_transf, batch_size=BATCH_SIZE, shuffle=False)

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0

for (bx, by), (bx2, by2) in zip(dl, dl_transf):
    for i, bxi in enumerate(bx):
        if spn == 0:
            print(bxi.shape)
        spn += 1
        if (spn > spr * spc):
            break
        plt.subplot(spr, spc, spn)
        bxi = bxi.transpose(0, 2)
        bxi = bxi.transpose(0, 1)
        plt.imshow(bxi)
        plt.axis('off')
        cls_id = by[i].item()
        plt.title(str(cls_id) + ': ' + mnist_train.classes[cls_id])

        spn += 1
        if (spn > spr * spc):
            break
        plt.subplot(spr, spc, spn)
        bxi = bx2[i]
        bxi = bxi.transpose(0, 2)
        bxi = bxi.transpose(0, 1)
        plt.imshow(bxi)
        plt.axis('off')
        cls_id = by2[i].item()
        plt.title(str(cls_id) + ': ' + mnist_train.classes[cls_id])
    if (spn > spr * spc):
        break
